/**
 * 
 */
package edu.umd.clip.lm.model.data;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.*;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class BufferedTrainingDataNodeWriter implements TrainingDataNodeWriter {
	private TrainingDataNodeWriter backend;
	private Lock lock;
	private Condition canGetSome;
	private Condition gotSome;
	private Queue<Action> pendingActions;
	// the size is in ContextFuturesPair records
	private long pendingDataSize;
	private long maxPendingDataSize;
	public static final long DEFAULT_MAX_PENDING_DATA_SIZE = 2 * 1024 * 1024;
	
	private AsyncWriterHelper writerThread;
	
	private static enum ActionType {
		ADD,
		START,
		FINISH,
		WRITE,
		QUIT
	}
	
	private static class Action {
		final ActionType type;
		final WritableTrainingDataNode node;
		final WritableTrainingData data;
		final TrainingDataBlock block;
		final long size;
		
		private Action(ActionType type, WritableTrainingData data, TrainingDataBlock block) {
			this(type, data, block, 0);
		}
		
		private Action(ActionType type, WritableTrainingData data, TrainingDataBlock block, long size) {
			this.node = null;
			this.type = type;
			this.data = data;
			this.block = block;
			this.size = size;
		}
		
		private Action(ActionType type, WritableTrainingDataNode node) {
			this.node = node;
			this.type = type;
			this.data = null;
			this.block = null;
			this.size = 0;
		}
		
		public static Action start(WritableTrainingData data) {
			return new Action(ActionType.START, data, null);
		}
		
		public static Action finish(WritableTrainingData data) {
			return new Action(ActionType.FINISH, data, null);
		}
		
		public static Action write(WritableTrainingData data, TrainingDataBlock block) {
			return new Action(ActionType.WRITE, data, block, block.size());
		}
		
		public static Action add(WritableTrainingDataNode node) {
			return new Action(ActionType.ADD, node);
		}
		
		public static Action quit() {
			return new Action(ActionType.QUIT, null);
		}
	}
	
	private class AsyncWriterHelper extends Thread {
		
		public AsyncWriterHelper() {
			super("async writer helper");
			setDaemon(true);
		}
		
		public void run() {
			boolean gotQuit = false;
			while(true) {
				lock.lock();
				Action action = null;
				try {
					while(pendingActions.isEmpty()) {
						try {
							gotSome.await();
						} catch (InterruptedException e) {}
					}
					action = pendingActions.poll();
				} finally {
					lock.unlock();
				}
				
				try {
					switch(action.type) {
					case ADD: 
						backend.add(action.node);
						break;
					case START:
						backend.start(action.data);
						break;
					case FINISH:
						backend.finish(action.data);
						break;
					case WRITE:
						backend.write(action.data, action.block);
						break;
					case QUIT:
						gotQuit = true;
						if (!pendingActions.isEmpty()) {
							System.err.printf("There are %d actions past QUIT. they will be dropped!!\n", pendingActions.size());
							pendingActions.clear();
						}
						if (pendingDataSize != 0) {
							System.err.printf("pendingDataSize=%d at QUIT, must be 0!\n", pendingDataSize);
							pendingDataSize = 0;
						}
						break;
					}
					try {
						lock.lock();
						pendingDataSize -= action.size;
						if (pendingDataSize < maxPendingDataSize) {
							canGetSome.signalAll();
						}
					} finally {
						lock.unlock();
					}
					
					if (gotQuit) return;
					
				} catch(IOException e) {
					e.printStackTrace();
				}
			}
		}
	}
	
	public BufferedTrainingDataNodeWriter(TrainingDataNodeWriter backend) {
		this(backend, DEFAULT_MAX_PENDING_DATA_SIZE);
	}

	public BufferedTrainingDataNodeWriter(TrainingDataNodeWriter backend, long bufferSize) {
		this.backend = backend;
		this.maxPendingDataSize = bufferSize;
		
		this.pendingActions = new LinkedList<Action>();
		this.pendingDataSize = 0;
		this.lock = new ReentrantLock();
		this.canGetSome = lock.newCondition();
		this.gotSome = lock.newCondition();
		
		this.writerThread = null;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.data.TrainingDataNodeWriter#add(edu.umd.clip.lm.model.data.WritableTrainingDataNode)
	 */
	@Override
	public void add(WritableTrainingDataNode node) {
		submitAction(Action.add(node));
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.data.TrainingDataNodeWriter#createNode(int, int)
	 */
	@Override
	public WritableTrainingDataNode createNode(int nodeId, int numData) {
		WritableTrainingData[] data = new WritableTrainingData[numData];
		for(int i=0; i<numData; ++i) {
			data[i] = createData();
		}
		return new WritableTrainingDataNode(nodeId, data);
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.data.TrainingDataWriter#createData()
	 */
	@Override
	public WritableTrainingData createData() {
		return new WritableTrainingData(this);
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.data.TrainingDataWriter#finish(edu.umd.clip.lm.model.data.WritableTrainingData)
	 */
	@Override
	public void finish(WritableTrainingData data) {
		submitAction(Action.finish(data));
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.data.TrainingDataWriter#start(edu.umd.clip.lm.model.data.WritableTrainingData)
	 */
	@Override
	public void start(WritableTrainingData data) {
		submitAction(Action.start(data));
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.data.TrainingDataWriter#write(edu.umd.clip.lm.model.data.WritableTrainingData, edu.umd.clip.lm.model.data.TrainingDataBlock)
	 */
	@Override
	public void write(WritableTrainingData data, TrainingDataBlock block) throws IOException {
		submitAction(Action.write(data, block));
	}

	private void submitAction(Action action) {
		try {
			lock.lock();
			// TODO: block is pendingDataSize >= maxPendingDataSize
			// TODO: writerThread handling is fishy (see sync() method)
			if (writerThread == null) {
				writerThread = new AsyncWriterHelper();
				writerThread.start();
			}
			pendingActions.add(action);
			pendingDataSize += action.size;
			gotSome.signalAll();
		} finally {
			lock.unlock();
		}
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.data.TrainingDataNodeWriter#sync()
	 */
	@Override
	public void sync() throws IOException {
		// TODO: it's not really sync() -- it's really a close()
		AsyncWriterHelper oldWriterThread = null;
		try {
			lock.lock();
			submitAction(Action.quit());
			/*
			while(!pendingActions.isEmpty()) {
				try {
					canGetSome.await(1, TimeUnit.SECONDS);
					gotSome.signal(); // awake the thread
				} catch (InterruptedException e) {}
			}
			*/
			
			if (writerThread != null) {
				//gotSome.signal(); // awake the thread
				//writerThread.interrupt();
				oldWriterThread = writerThread;
				writerThread = null;
			}
		} finally {
			lock.unlock();
		}

		// wait for the thread to finish the last operation
		try {
			if (oldWriterThread != null) {
				oldWriterThread.join();
			}
		} catch (InterruptedException e) {
			e.printStackTrace();
		}
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.data.TrainingDataNodeWriter#close()
	 */
	@Override
	public void close() throws IOException {
		sync();
		backend.close();
	}
}
